#!/usr/bin/env perl
# by lhtk : lhtk80t7@gmail.com
# 
# 自动化流程，为 B2DSC 添加数据。
# 
use strict;
use warnings;
use JSON;
use Cwd 'abs_path';
use Getopt::Long;

my $fa;
my $chrom_id_tab; # Id 对应关系表
my $gz;
my $bgz;
my $help;
my $b2dsc_dir;
my $dbname;
my $rm_raw_fasta; # 是否删除原始参考基因组序列文件

GetOptions(
    'fasta|f=s' => \$fa,
    'table|t=s' => \$chrom_id_tab,
    'gzip|z'    => \$gz,     # flag
    'bgzip|bgz' => \$bgz,    # flag
    'help|h'    => \$help,
    'b2dsc_dir|b2dsc|b=s'   => \$b2dsc_dir,
    'dbname|n=s'  => \$dbname,
    'remove_raw_fasta|r' => \$rm_raw_fasta, # flag
) or die "Error in command line arguments\n";

my %dirs;

check_args();
check_dependency(); # 软件依赖

# ===============================================================
# 1. cid    hgrp    sub-genome
# 
my %hs; # hgrp, subG
my %cids;
print STDERR "读取 $chrom_id_tab\n";
open my $fh_tab, '<', $chrom_id_tab or die $!;
while (<$fh_tab>) {
    chomp;
    if (/^(\S+)\t(\d+)\t([a-zA-Z]\w*)\b/) {
        my($id,$hgrp,$subG) = ($1,$2,$3);
        $hs{$id} = {
            hgrp => $hgrp, subG => $subG,
        };
        $cids{$subG}{$hgrp} = $id;
    }
}
close $fh_tab or die $!;
die "未从染色体对应关系表中获得任何信息！\n" unless %hs;
for my $subG (keys %cids) {
    my @hgrps = sort { $a <=> $b } keys %{ $cids{$subG} };
    die "同源群必须从 1 开始！\n" unless $hgrps[0] == 1;
    for my $i (1 .. $#hgrps) {
        die "同源群必须为从 1 开始的一系列正整数（如 1, 2, 3, ...），不可跳漏！\n" if $hgrps[$i] != $hgrps[$i - 1] + 1;
    }
}

# ===============================================================
# 2. 表没问题再新建文件夹
# 
print STDERR "创建文件夹\n";
for my $dn (qw/bdb Ns fasta/) {
    system "mkdir -p $dirs{$dn}";
}

# ===============================================================
# 3. bgzip 文件，构建索引
# 
print STDERR "准备 bgzip 压缩的参考基因组 FASTA\n";
# my $fa_bgz;
chomp(my $fa_bname = `basename $fa .gz`); # basename，类似 xxx.fa
my $fa_bgz = "$dirs{fasta}/$fa_bname.gz";
die "$fa_bgz: 已存在！" if -e $fa_bgz;
unless ($bgz) {
    if ($gz) {
        system "zcat $fa | bgzip > $fa_bgz";
    } else {
        system "cat $fa | bgzip > $fa_bgz";
    }
} else {
    system "cp -i $fa $dirs{fasta}/";
}
$fa_bname = "$fa_bname.gz";
system "samtools faidx $fa_bgz";
my $fai = "$fa_bgz.fai";
# ===============================================================
# 4. 构建 blastdb
# 
print STDERR "构建 blastdb\n";
open my $fh_fai, '<', $fai or die $!;
my %clen; # chromosome length
while (<$fh_fai>) {
    next if /^\s*$/;
    my @r = split /\t/;
    next if @r < 2;
    my($id,$bp) = @r[0,1];
    next if not exists $hs{$id};
    $clen{$id} = $bp;
    make_bdb($id, $hs{$id}{hgrp} . $hs{$id}{subG});
}
close $fh_fai or die $!;
for my $id (keys %hs) {
    die "提供的 $chrom_id_tab 中的序列 $id 未见于 FASTA 文件 $fa！\n" unless exists $clen{$id};
}

# ===============================================================
# 5. Ns
# 
print STDERR "寻找参考基因组中的未知碱基 N 的信息\n";
find_Ns_samtools();

# ===============================================================
# 6. conf
# 
my %conf_b2dsc;
my %conf_faidx;
$conf_b2dsc{blastdb}{db_dir} = $dirs{bdb};
$conf_b2dsc{N}{dir} = $dirs{Ns};
$conf_faidx{dir} = $dirs{fasta};
for my $subG (keys %cids) {
    my @lens = ();
    for my $hgrp (sort { $a <=> $b } keys %{ $cids{$subG} }) {
        my $did = $hgrp . $subG;
        my $cid = $cids{$subG}{$hgrp};
        $conf_b2dsc{blastdb}{chromosome}{$subG}{$hgrp} = [$did, $cid];
        $conf_b2dsc{N}{chromosome}{$subG}{$hgrp} = "$did.tsv";
        $conf_faidx{chromosome}{$subG}{$hgrp} = [
            $fa_bname, $cid, $clen{$cid}
        ];
        push @lens, $clen{$cid};
    }
    $conf_b2dsc{chromosome_length}{$subG} = \@lens;
}
print STDERR "写入配置文件\n";
my $json = JSON->new->allow_nonref;
my $json_b2dsc = join('/', $dirs{conf}{b2dsc}, "${dbname}.json");
my $json_faidx = join('/', $dirs{conf}{faidx}, "${dbname}.json");
my $fh_w;
open $fh_w, '>', $json_b2dsc or die $!;
print $fh_w $json->pretty->encode(\%conf_b2dsc);
close $fh_w or die $!;

open $fh_w, '>', $json_faidx or die $!;
print $fh_w $json->pretty->encode(\%conf_faidx);
close $fh_w or die $!;

# ===============================================================
# 7. 删除原始的 FASTA
# 
if ($rm_raw_fasta) {
    print STDERR "删除原始的 FASTA\n";
    unlink $fa;
}

print STDERR "程序结束！\n";

# ===============================================================
# subroutines
# ===============================================================

sub check_args {
    usage() if ($help);

    unless (defined $chrom_id_tab and -T $chrom_id_tab) {
        die "必须指定染色体编号对应关系表！\n";
    }

    unless (defined $fa and ($gz or $bgz or -T $fa)) {
        die "必须提供 fasta 文件！可能你提供的是 gzip/bgzip 压缩文件，而未启用 -z 或 --bgz 选项！\n";
    }

    if ($gz or $bgz) {
        die "$fa: 后缀必须是 --> .gz <--！\n" unless ($fa =~ /\.gz$/);
    }

    unless (defined $dbname and $dbname =~ /^[\w\-]+$/) {
        die "请设定正确的数据库名称！（-n）\n";
    }

    unless (defined $b2dsc_dir) {
        # $b2dsc_dir = '/home/lhtk/sites/mcgb';
        die "必须指定 B2DSC 安装目录，类似: /home/lhtk/git/b2dsc\n";
    } else {
        if ($b2dsc_dir =~ /^~/) { # e.g. ~/git/b2dsc
            $b2dsc_dir = glob($b2dsc_dir);
        } else {
            $b2dsc_dir = abs_path($b2dsc_dir);
        }
    }
    die "$b2dsc_dir: 该文件夹不存在！\n" unless -d $b2dsc_dir;

    if (not -o $b2dsc_dir) {
        die "当前用户并非文件夹 --> $b2dsc_dir <-- 的所有者！\n请切换至正确账户!\n";
    }

    %dirs = (
        bdb => join('/', $b2dsc_dir, 'data/blastdb', $dbname),
        Ns => join('/', $b2dsc_dir, 'data/Ns', $dbname),
        fasta => join('/', $b2dsc_dir, 'data/fasta', $dbname),
        conf => {
            b2dsc => join('/', $b2dsc_dir, 'conf/b2dsc/species'),
            faidx => join('/', $b2dsc_dir, 'conf/faidx'),
        },
    );

    for my $cdir (qw/b2dsc faidx/) {
        if (-e $dirs{conf}{$cdir} . '/' . "${dbname}.json") {
            die "已有同名数据库，请重新为 -n 选项赋值！\n";
        }
    }
    
    for my $dn (qw/bdb Ns fasta/) {
        if (-e $dirs{$dn}) {
            die "已有同名数据库，请重新为 -n 选项赋值！\n";
        }
    }
}

sub usage {
    print <<USAGE;

为 B2DSC 准备数据:
    1. bgzip 文件
    2. blastdb
    3. Ns
    4. 配置文件
依赖：
    1. NCBI-BLAST+ 的 makeblastdb
    2. samtools
    3. bgzip

用法示例：
    # 若是未压缩的纯文本 FASTA 文件：
    $0 -f IWGSC_RefSeq_V1_chr1A.fsa -n wheat_1A -r -t chrom-id-table.tab -b ~/git/b2dsc
    #
    # 注意：-r 是删除用户提供的原始 FASTA 文件，若不删除，则勿启用该选项。

    # 若是普通 gzip 文件：
    $0 -f IWGSC_RefSeq_V1_chr1A.fsa.gz -z -n wheat_1A -r -t chrom-id-table.tab -b /home/lhtk/git/b2dsc

    # 若是 bgzip 构建的 gzip 文件：
    $0 -f IWGSC_RefSeq_V1_chr1A.fsa.gz -bgz -n wheat_1A -t chrom-id-table.tab -b /home/lhtk/sites/mcgb

选项：

    --fasta, -f
        要分析的 FASTA 文件，必须是 DNA 序列。

    --table, -t
        染色体 ID 对应关系表，包含 3 列数据，以 Tab 分隔。
        第 1 列: FASTA 文件中的原始序列编号。
        第 2 列: 序列所属同源群，必须为从 1 开始的一系列正整数。
        第 3 列: 序列所属亚基因组，必须为英文字母、数字、下划线组成，必须以英文字母开头。

    --gzip, -z
        fasta 文件是 gzip 格式的。
        可选。
        flag

    --bgzip, --bgz
        提供了 bgzip 压缩后的 fasta 文件。
        可选。
        flag

    --b2dsc_dir, --b2dsc, -b
        B2DSC 或者 MCGB 网站所在目录。
        例如: /home/lhtk/sites/mcgb 或 /home/lhtk/git/b2dsc

        注意：必须使用这一目录所有者账号，例如 lhtk 用户。

    --dbname, -n
        数据库名称，请仅使用以下字符：
        1. 英文字母 a-zA-Z
        2. 阿拉伯数字 0-9
        3. 下划线 _
        4. 短横线 -

        例如: Aegilops_speltoides

    --remove_raw_fasta, -r
        是否删除原始的参考基因组 FASTA 文件？
        可选。
        flag

    --help, -h
        打印此帮助。
        可选。
        flag
USAGE
    exit 0;
}

sub make_bdb {
    my($cid,$cname) = @_;
    system "samtools faidx $fa_bgz $cid | makeblastdb -in - -dbtype nucl -title '$dbname $cid' -parse_seqids -out $dirs{bdb}/$cname";
}

sub find_Ns_samtools {
    my $out_dir = $dirs{Ns};

    my @seq_list = keys %clen; # made bdb
    for my $sid (@seq_list) {
        my $p = 0;    # pos
        my $s;        # start

        open my $fh, "samtools faidx -n 60 $fa_bgz $sid |" or die $!;
        my $f = join '/', $out_dir, $hs{$sid}{hgrp} . $hs{$sid}{subG} . ".tsv";
        print STDERR "输出到 $f ...\n";
        open my $fh_o, '>', $f or die $!;

        while (<$fh>) {
            next if /^>/;
            process_line( $_, $fh_o, \$s, \$p );
        }
        print $fh_o "$p\n" if defined $s;
        close $fh   or die $!;
        close $fh_o or die $!;
    }
}

sub process_line {
    my ( $line, $fh_o, $s, $p ) = @_;

    return if ( $line =~ /^\s*$/ );
    $line =~ s/\s//g;

    # 注意一行全是 N，但下一行开始不是 N 了
    unless ( $line =~ /N/i ) {
        if ( defined $$s ) {
            print $fh_o $$p, "\n";
            $$s = undef;
        }
        $$p += length $line;
        return;
    }

    for my $nc ( split //, $line ) {
        $$p++;
        if ( $nc !~ /N/i ) {
            if ( defined $$s ) {
                my $e = $$p - 1;
                print $fh_o $e . "\n";
                $$s = undef;
            }
        }
        else {
            if ( not defined $$s ) {
                $$s = $$p;
                print $fh_o "$$s\t";
            }
        }
    }
}

sub check_dependency {
    for my $exe (qw/makeblastdb samtools bgzip/) {
        my $res = `which $exe`;
        die "$exe 程序不可用，请将其文件夹添加到环境变量 PATH 中！\n" unless ( $res =~ /\S+/ );
    }
}